import concurrent.futures
from typing import List, Callable, TypeVar, Generic
import multiprocessing
from functools import partial
import logging

T = TypeVar('T')
R = TypeVar('R')

class ParallelProcessor(Generic[T, R]):
    def __init__(self, max_workers: int = None):
        """
        初始化并行处理器
        :param max_workers: 最大工作线程数，默认为CPU核心数
        """
        self.max_workers = max_workers or multiprocessing.cpu_count()
        self.logger = logging.getLogger(__name__)

    def process(self, items: List[T], process_func: Callable[[T], R]) -> List[R]:
        """
        并行处理项目列表
        :param items: 要处理的项目列表
        :param process_func: 处理函数
        :return: 处理结果列表
        """
        results = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有任务
            future_to_item = {
                executor.submit(process_func, item): item
                for item in items
            }

            # 收集结果
            for future in concurrent.futures.as_completed(future_to_item):
                item = future_to_item[future]
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    self.logger.error(f"处理项目时发生错误: {str(e)}")
                    raise

        return results

    def process_with_progress(self, items: List[T], process_func: Callable[[T], R],
                            progress_callback: Callable[[int, int], None] = None) -> List[R]:
        """
        带进度回调的并行处理
        :param items: 要处理的项目列表
        :param process_func: 处理函数
        :param progress_callback: 进度回调函数，接收当前进度和总数
        :return: 处理结果列表
        """
        results = []
        total_items = len(items)
        completed = 0

        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有任务
            future_to_item = {
                executor.submit(process_func, item): item
                for item in items
            }

            # 收集结果
            for future in concurrent.futures.as_completed(future_to_item):
                item = future_to_item[future]
                try:
                    result = future.result()
                    results.append(result)
                    completed += 1
                    if progress_callback:
                        progress_callback(completed, total_items)
                except Exception as e:
                    self.logger.error(f"处理项目时发生错误: {str(e)}")
                    raise

        return results

    def process_batch(self, items: List[T], process_func: Callable[[List[T]], List[R]],
                     batch_size: int = 10) -> List[R]:
        """
        批量并行处理
        :param items: 要处理的项目列表
        :param process_func: 批量处理函数
        :param batch_size: 批处理大小
        :return: 处理结果列表
        """
        # 将项目分成批次
        batches = [items[i:i + batch_size] for i in range(0, len(items), batch_size)]
        
        results = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # 提交所有批次任务
            future_to_batch = {
                executor.submit(process_func, batch): batch
                for batch in batches
            }

            # 收集结果
            for future in concurrent.futures.as_completed(future_to_batch):
                batch = future_to_batch[future]
                try:
                    batch_results = future.result()
                    results.extend(batch_results)
                except Exception as e:
                    self.logger.error(f"处理批次时发生错误: {str(e)}")
                    raise

        return results 